import sys
from PyQt5.QtWidgets import *
from PyQt5 import uic, QtWidgets, QtGui
import skimage
from skimage.metrics import structural_similarity as ssim
from merge_image_core import Our_Image, read_h5ad,aligning_img, set_binary,aligning_img_ori
from adjust_window import SetWindow
from merge_window import Merge_Window
import h5py
import cv2
import pickle
import os
import numpy as np
import tifffile

def resource_path(relative_path):
    import sys, os
    if hasattr(sys, '_MEIPASS'):
        return os.path.join(sys._MEIPASS, relative_path)
    return os.path.abspath(relative_path)

class MainWindow(QWidget):

    def __init__(self):
        super().__init__()
        self.init_ui()

    def init_ui(self):
        self.ui = uic.loadUi(resource_path("gene_page.ui"))
        self.button_he = self.ui.pushButton_2
        self.button_fluo = self.ui.pushButton_3
        self.button_gene = self.ui.pushButton_4
        self.grap_he = self.ui.graphicsView_HE
        self.grap_fluo = self.ui.graphicsView_FLUO
        self.grap_gene = self.ui.graphicsView_GENE
        self.button_he_ad = self.ui.pushButton_he
        self.button_fluo_ad = self.ui.pushButton_fluo
        self.button_gene_ad = self.ui.pushButton_gene

        self.button_he_merge = self.ui.pushButton_He
        self.button_fluo_merge = self.ui.pushButton_Fluo
        self.button_gene_merge  = self.ui.pushButton_Gene
        self.grap_merge = self.ui.graphicsView
        self.button_end = self.ui.pushButton

        self.button_score_he = self.ui.pushButton_5
        self.button_score_he.clicked.connect(self.caculate_score_he)
        self.SSIM_score_he = self.ui.label_4

        self.button_score_fluo = self.ui.pushButton_6
        self.button_score_fluo.clicked.connect(self.caculate_score_fluo)
        self.SSIM_score_fluo = self.ui.label_6

        self.img_he = None #image_treat(img=None, x=320, y=320)
        self.img_fluo = None # image_treat(img=None, x=240, y=240)
        self.img_gene = None #image_treat(img=None, x=240, y=240)
        self.obj_merge = None

        #信号和槽传递参数
        #加载图片
        self.button_he.clicked.connect(lambda: self.loadimage(name = 'he'))
        self.button_fluo.clicked.connect(lambda: self.loadimage(name = 'fluo'))
        self.button_gene.clicked.connect(lambda: self.loadimage(name = 'gene'))

        #调整图片
        self.button_he_ad.clicked.connect(lambda: self.adjustimage(name='he'))
        self.button_fluo_ad.clicked.connect(lambda: self.adjustimage(name='fluo'))
        self.button_gene_ad.clicked.connect(lambda: self.adjustimage(name='gene'))

        #合并图片
        self.button_he_merge.clicked.connect(lambda: self.mergeimage(name='he'))
        self.button_fluo_merge.clicked.connect(lambda: self.mergeimage(name='fluo'))
        self.button_gene_merge.clicked.connect(lambda: self.mergeimage(name='gene'))

        self.scene_view = QGraphicsScene()
        self.scene_view.setSceneRect(0, 0, 860, 860)
        brush = QtGui.QBrush(QtGui.QColor(0, 0, 0))
        #brush.setStyle(QtCore.Qt.SolidPattern)
        # self.graphicsView.setForegroundBrush(brush)
        #self.scene_view.setBackgroundBrush(brush)

        #self.sig_merge_draw.connect(self.merge_draw)
        self.button_end.clicked.connect(self.save_end)

    def loadimage(self,name):
        if name == 'he':
            fname, _ = QFileDialog.getOpenFileName(self, 'Open file he', '..')
            img = skimage.io.imread(fname)
            print("HE shape:", img.shape)
            self.img_he = Our_Image(img)
            self.img_he.draw_img(self.grap_he)

        if name == 'fluo':
            fname, _ = QFileDialog.getOpenFileName(self, 'Open file fluo', '..')
            img = skimage.io.imread(fname)
            self.img_fluo = Our_Image(img)
            print("FLUO shape:", img.shape)
            self.img_fluo.draw_img(self.grap_fluo)

        if name == 'gene':
            fname, _ = QFileDialog.getOpenFileName(self, 'Open file gene', '..')
            img = read_h5ad(fname)
            self.h5ad_path = fname
            self.img_gene = Our_Image(img)
            print("GENE shape:", img.shape)
            self.img_gene.draw_img(self.grap_gene)

    def adjustimage(self, name):
        if name == 'he':
            adjustWidget = SetWindow(self.img_he, grap=self.grap_he, name = 'he')
            adjustWidget.ui.show()
            self.img_he.draw_img(self.grap_he)

        if name == 'fluo':
            adjustWidget = SetWindow(self.img_fluo, grap=self.grap_fluo, name = 'fluo')
            adjustWidget.ui.show()
            self.img_fluo.draw_img(self.grap_fluo)

        if name == 'gene':
            adjustWidget = SetWindow(self.img_gene, self.grap_gene, name = 'gene')
            adjustWidget.ui.show()
            self.img_gene.draw_img(self.grap_gene)

    def mergeimage(self, name):
        if name == 'he':
            self.mergeWidget_he = Merge_Window(self.img_he, self.grap_merge, self.scene_view,
                                               name = 'he')
            self.mergeWidget_he.ui.show()

        if name == 'fluo':
            self.mergeWidget_fluo = Merge_Window(self.img_fluo,self.grap_merge,self.scene_view,
                                            name = 'fluo')
            self.mergeWidget_fluo.ui.show()

        if name == 'gene':
            self.mergeWidget_gene = Merge_Window(self.img_gene,self.grap_merge,self.scene_view,
                                             name = 'gene')
            self.mergeWidget_gene.ui.show()

    def save_end(self):
        fname, _ = QFileDialog.getSaveFileName(self, 'Open file ', '../merge.h5')

        self.img_fluo.remove_exp_preserve()
        self.img_he.remove_exp_preserve()
        self.img_gene.remove_exp_preserve()

        img_gene_fluo = aligning_img(self.img_gene,self.img_fluo)
        img_gene_he = aligning_img(self.img_gene,self.img_he)
        img_gene_fluo_ori,scale_ori = aligning_img_ori(self.img_gene,self.img_fluo)
        img_gene_he_ori,scale_ori_he = aligning_img_ori(self.img_gene,self.img_he)
    
        # 校正图片channels
        # channel_he = channel_he = np.argmax(np.sum(np.sum(img_gene_he, axis=0), axis=0))
        # print("channel_he",channel_he)
        # if  channel_he != 0:
        #     zeros_channels = [i for i in [0,1,2] if i != channel_he]
        #     img_gene_he = img_gene_he[:,:,[channel_he,zeros_channels[0], zeros_channels[1]]]
        #     img_gene_he_ori = img_gene_he_ori[:,:,[channel_he,zeros_channels[0], zeros_channels[1]]]

        print(img_gene_fluo_ori.shape)
        print(scale_ori)
        print(img_gene_he_ori.shape)
        print(scale_ori_he)

        fluo_fp =  os.path.join(os.path.dirname(fname), 'fluo_ori.tiff')
        tifffile.imwrite(fluo_fp, img_gene_fluo_ori.astype(np.uint8),  bigtiff=True)

        fluo_fp =  os.path.join(os.path.dirname(fname), 'fluo_scale.tiff')
        tifffile.imwrite(fluo_fp, img_gene_fluo)

        he_fp =  os.path.join(os.path.dirname(fname), 'he_ori.tiff')
        tifffile.imwrite(he_fp, img_gene_he_ori.astype(np.uint8),  bigtiff=True)
        print("img_gene_he_ori.shape, img_gene_he_ori.dtype")
        print(img_gene_he_ori.shape, img_gene_he_ori.dtype)

        he_fp =  os.path.join(os.path.dirname(fname), 'he_scale.tiff')
        tifffile.imwrite(he_fp, img_gene_he)

        if os.path.exists(fname):
            os.remove(fname)

        file = h5py.File(fname,'w')

        #img_gene_fluo = set_binary(img_gene_fluo, self.img_fluo.a, self.img_fluo.b, self.img_fluo.c)
        #img_gene_he = set_binary(img_gene_he,self.img_he.a, self.img_he.b,self.img_he.c)

        file.create_dataset('he',data = img_gene_he)
        file.create_dataset('he-ori',data= img_gene_he_ori)
        file.create_dataset('he_scale',data=scale_ori_he)

        file.create_dataset('fluo', data =img_gene_fluo)
        file.create_dataset('fluo-ori',data = img_gene_fluo_ori)
        file.create_dataset('fluo_scale',data = scale_ori)
        file.close()

        # save_parameter
        self.img_gene.remove_images()
        self.img_fluo.remove_images()
        self.img_he.remove_images()
        new_fname_gene = fname.split('.')[0] + '_img_gene.pkl'
        new_fname_fluo = fname.split('.')[0] + '_img_fluo.pkl'
        new_fname_he = fname.split('.')[0] + '_img_he.pkl'

        with open(new_fname_gene, 'wb') as f:
            pickle.dump(self.img_gene, f)
        with open(new_fname_fluo, 'wb') as f:
            pickle.dump(self.img_fluo, f)
        with open(new_fname_he, 'wb') as f:
            pickle.dump(self.img_he, f)

        #data = h5_to_h5ad(self.h5ad_path, img_gene_he, img_gene_fluo)
        #QtWidgets.QMessageBox.question(self, "Xpath Robot", "保存成功",
        # QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No)
        QtWidgets.QMessageBox.question(self, "Xpath Robot", "保存成功")

    def caculate_score_he(self):
        img_he = aligning_img(self.img_gene, self.img_he)
        img_he = set_binary(img_he, self.img_he.a, self.img_he.b, self.img_he.c)
        img_he = cv2.cvtColor(img_he,cv2.COLOR_RGB2GRAY)
        img_gene = self.img_gene.img_preserve
        img_gene = cv2.normalize(img_gene, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        score = ssim(img_gene,img_he)
        self.SSIM_score_he.setText('%.8f'%(score))

    def caculate_score_fluo(self):
        img_fluo = aligning_img(self.img_gene, self.img_fluo)
        img_fluo = set_binary(img_fluo, self.img_fluo.a, self.img_fluo.b, self.img_fluo.c)
        img_fluo = cv2.cvtColor(img_fluo,cv2.COLOR_RGB2GRAY)
        img_gene = self.img_gene.img_preserve
        img_gene = cv2.normalize(img_gene, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        score = ssim(img_gene,img_fluo)
        self.SSIM_score_fluo.setText('%.8f'%(score))

if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = MainWindow()
    w.ui.show()
    app.exec()

